#include "Node/LRG_Node_ComponentBuckets.h"

#include "Component/LineReplicationGraphComponent.h"
#include "Component/LRG_ViewerComponent.h"

ULRG_Node_ComponentBuckets::ULRG_Node_ComponentBuckets()
{
	SetNonStreamingCollectionSize(Settings.NumBuckets);
}

void ULRG_Node_ComponentBuckets::SetNonStreamingCollectionSize(int32 NewSize)
{
	NonStreamingComponentCollection.SetNum(Settings.NumBuckets);
	//static TArray<FLineComponentContainer*> FullList;
	//FullList.Reset();
	//for (FLineComponentContainer& List : NonStreamingComponentCollection)
	//{
	//	List.AppendToTArray(FullList);
	//}
	//// Readd/Rebalance
	//for (int32 idx=0; idx < FullList.Num(); ++idx)
	//{
	//	NonStreamingCollection[idx % NewSize].Add( FullList[idx] );
	//}
}

void ULRG_Node_ComponentBuckets::AddNetworkActor(ULineReplicationGraphComponent* Component)
{
	if(!Component)
	{
		return;
	}
	const FName& StreamingLevelName = Component->GetNewReplicatedActorInfo().StreamingLevelName;
	if(!StreamingLevelName.IsNone())
	{
		LevelStreamingComponentCollection.AddActor(Component);
		return;
	}

	FLineComponentContainer* BestContainer = nullptr;
	int32 LeastNum = INT_MAX;
	for(FLineComponentContainer& Value : NonStreamingComponentCollection)
	{
		if(LeastNum > Value.Num())
		{
			LeastNum = Value.Num();
			BestContainer = &Value;
		}
	}
	BestContainer->Add(Component);
}

bool ULRG_Node_ComponentBuckets::RemoveNetworkActor(ULineReplicationGraphComponent* Component)
{
	if(!Component)
	{
		return false;
	}
	const FName& StreamingLevelName = Component->GetNewReplicatedActorInfo().StreamingLevelName;
	if(!StreamingLevelName.IsNone())
	{
		return LevelStreamingComponentCollection.RemoveActor(Component);
	}
	bool bFound = false;
	for(FLineComponentContainer& Value : NonStreamingComponentCollection)
	{
		if(Value.Remove(Component))
		{
			bFound = true;
		}
	}
	return bFound;
}

void ULRG_Node_ComponentBuckets::ResetAllNetworkActors()
{
	for(FLineComponentContainer& Value : NonStreamingComponentCollection)
	{
		Value.Reset();
	}
	LevelStreamingComponentCollection.Reset();
}

void ULRG_Node_ComponentBuckets::GatherActorListsForConnection(const FConnectionGatherActorListParameters& Params)
{
	const uint32 ReplicationFrameNum = Params.ReplicationFrameNum;
	TArray<FName, FReplicationGraphConnectionsAllocator> ViewerLines;
	for (const FNetViewer& CurViewer : Params.Viewers)
	{
		if(ULRG_ViewerComponent* Component = ULRG_ViewerComponent::Get(CurViewer.ViewTarget) )
		{
			const TSet<FName>& VisibleLines = Component->GetVisibleLines();
			for (const FName& Line : VisibleLines)
			{
				ViewerLines.AddUnique(Line);
			}
		}
	}
	// Return one list as Default and the rest as FastShared
	const int32 DefaultReplicationIdx = (ReplicationFrameNum % NonStreamingComponentCollection.Num());
	for (int32 idx = 0; idx < NonStreamingComponentCollection.Num(); ++idx)
	{
		ListForGather.Reset();
		const FLineComponentContainer& Container = NonStreamingComponentCollection[idx];
		for(const auto& Component : Container.Components)
		{
			if(ViewerLines.Contains(Component->GetCurrentLine()))
			{
				ListForGather.Add(Component->GetOwner());
			}
		}
		if (DefaultReplicationIdx == idx)
		{
			// Default Rep Path
			Params.OutGatheredReplicationLists.AddReplicationActorList(ListForGather, EActorRepListTypeFlags::Default);
		}
		else
		{
			// Only do FastShared if modulo passes
			if (Settings.EnableFastPath && ReplicationFrameNum % Settings.FastPathFrameModulo == 0)
			{
				Params.OutGatheredReplicationLists.AddReplicationActorList(ListForGather, EActorRepListTypeFlags::FastShared);
			}
		}
	}

	LevelStreamingComponentCollection.Gather(Params);
}